!***************************************************************************************************************
! mUniverse contains all the cosmological parameters we need to describe a universe
!
! This is a base type for many other objects. There are more cosmological parameters in type(tLPS), for example.
!
! Declare the object:
!
! type(tUniverse) :: universe
!
! and initialize it with new:
!
! call new(universe, [Omega_m0], [Omega_v0], [w_0], [w_a], [Omega_b0], [h], [A_s], [sigma_8], [b_1], [b_2]
!
! This is your handle on all the cosmological parameters. You can declare multiple universe objects with different
! parameter combinations. All parameters are optional; the defaults are in base/parameters.F90.
!
! The growth factor is here (and normalized automatically by new):
!
! growthFactor(universe, a)
!
! Also some handy conversion functions: m_R(universe, m), R_m(universe, R), to convert mass to Lagrangian radius,
! and z_a and a_z for redshift to scale factor.
!
! The Omegas for new are capital Omegas. You can get and set the little omegas, omega_X = Omega_X * h^2 using
!
! getCMBParameters(universe, omega_b, omega_d, Omega_DE) and
! setCMBParameters(universe, omega_b, omega_d, Omega_DE)
!
!***************************************************************************************************************

module mUniverse

        use mPrecision
        use mObject

        implicit none

        real(D), parameter :: AVG_DENSITY_OVER_H2_OMEGA_M0 = 2.7755e11_D

        real(D), parameter :: SOURCE_SCALE_FACTOR = 0.5_D ! z = 1

        interface new
                module procedure new_universe
        end interface new

        interface write
                module procedure write_universe
        end interface write

        interface getH
                module procedure getH_universe
        end interface getH

        interface getOmega_m0
                module procedure getOmega_m0_universe
        end interface getOmega_m0

        interface getOmega_v0
                module procedure getOmega_v0_universe
        end interface getOmega_v0

        interface getOmega_b0
                module procedure getOmega_b0_universe
        end interface getOmega_b0

        interface getOmega_b
                module procedure getOmega_b_universe
        end interface getOmega_b

        interface setOmega_b
                module procedure setOmega_b_universe
        end interface setOmega_b

        interface getOmega_d
                module procedure getOmega_d_universe
        end interface getOmega_d

        interface setOmega_d
                module procedure setOmega_d_universe
        end interface setOmega_d

        interface getOmega_de
                module procedure getOmega_de_universe
        end interface getOmega_de

        interface setOmega_de
                module procedure setOmega_de_universe
        end interface setOmega_de

        interface omegas
                module procedure omegas_universe
        end interface omegas

        interface getW_0
                module procedure getW_0_universe
        end interface getW_0

        interface setW_0
                module procedure setW_0_universe
        end interface setW_0

        interface getW_a
                module procedure getW_a_universe
        end interface getW_a

        interface setW_a
                module procedure setW_a_universe
        end interface setW_a

        interface getB_1
                module procedure getB_1_universe
        end interface getB_1

        interface setB_1
                module procedure setB_1_universe
        end interface setB_1

        interface getB_2
                module procedure getB_2_universe
        end interface getB_2

        interface setB_2
                module procedure setB_2_universe
        end interface setB_2

        interface getShapeParameter
                module procedure getShapeParameter_universe
        end interface getShapeParameter

        interface getAvgDensity
                module procedure getAvgDensity_universe
        end interface getAvgDensity

        interface getA_s
                module procedure getA_s_universe
        end interface getA_s

        interface getSigma_8
                module procedure getSigma_8_universe
        end interface getSigma_8

        private
        public :: new, write
        public :: growthFactor, exactGrowthFactor
        public :: getOmega_m0, getOmega_v0, getOmega_b0, getSigma_8, getShapeParameter
        public :: omegas
        public :: getH
        public :: getW_0, setW_0
        public :: getW_a, setW_a
        public :: setB_1, setB_2, getB_1, getB_2
        public :: getA_s
        public :: getCMBParameters, setCMBParameters
        public :: getOmega_b, setOmega_b
        public :: getOmega_d, setOmega_d
        public :: getOmega_de, setOmega_de
        public :: getAvgDensity
        public :: equal
        public :: SOURCE_SCALE_FACTOR
        public :: hubble, m_R, R_m
        public :: a_z, z_a

contains

        subroutine new_universe(u, omega_m0, omega_v0, w_0, w_a, omega_b0, h, A_s, sigma_8, b_1, b_2)

                use mParameters

                type(tUniverse), intent(out) :: u
                real(D), intent(in), optional :: omega_m0, omega_v0, w_0, w_a, omega_b0, h, A_s, sigma_8, b_1, b_2

                if( present(omega_m0) ) then
                        u%omega_m0 = omega_m0
                else
                        u%omega_m0 = DEFAULT_OMEGA_M0
                end if

                if( present(omega_v0) ) then
                        u%omega_v0 = omega_v0
                else
                        u%omega_v0 = DEFAULT_OMEGA_V0
                end if

                if( present(w_0) ) then
                        u%w_0 = w_0
                else
                        u%w_0 = DEFAULT_W_0
                end if

                if( present(w_a) ) then
                        u%w_a = w_a
                else
                        u%w_a = DEFAULT_W_A
                end if

                if( present(omega_b0) ) then
                        u%omega_b0 = omega_b0
                else
                        u%omega_b0 = DEFAULT_OMEGA_B0
                end if

                if( present(h) ) then
                        u%h = h
                else
                        u%h = DEFAULT_H
                end if

                if( present(A_s) ) then
                        u%A_s = A_s
                else
                        u%A_s = DEFAULT_A_S
                end if

                if( present(sigma_8) ) then
                        u%sigma_8 = sigma_8
                else
                        u%sigma_8 = DEFAULT_SIGMA_8
                end if

                if( present(b_1) ) then
                        u%b_1 = b_1
                else
                        u%b_1 = DEFAULT_B_1
                end if

                if( present(b_2) ) then
                        u%b_2 = b_2
                else
                        u%b_2 = DEFAULT_B_2
                end if

!                u%shapeParameter = u%omega_m0 * u%h

!                u%avgDensity = AVG_DENSITY_OVER_H2_OMEGA_M0 * u%h * u%h * u%omega_m0

                call normalize(u)

        end subroutine new_universe

        subroutine normalize(u)

                type(tUniverse), intent(inout) :: u
                
                u%growthFactorNorm = 1.0_D
                u%growthFactorNorm = 1.0_D / growthFactor(u, 1.0_D)

        end subroutine normalize

        !**********************************************
        ! This is the approximate formula from Carroll:
        !**********************************************
        real(D) function growthFactor(u, a)

                use mDiffEq

                type(tUniverse), intent(in) :: u
                real(D), intent(in) :: a

                real(D), parameter :: A_BOUNDARY_CONDITIONS = 1.0_D / 1100.0_D ! i.e., decoupling
                real(D), parameter :: BOUNDARY_CONDITIONS(2) = (/ 1.0_D, 0.0_D /)
                real(D), parameter :: EPS = 0.009_D

                real(D) :: omega_m, omega_v
                type(tDiffEq) :: growthFactorEqs

                if( darkEnergy(u) ) then
                        call new(growthFactorEqs, A_BOUNDARY_CONDITIONS, BOUNDARY_CONDITIONS, eps = EPS, param = (/ u%omega_m0, u%w_0, u%w_a /) )
                        call integrate(growthFactorEqs, a, darkEnergyGrowthFactorDerivs)
                        growthFactor = u%growthFactorNorm * a * y(growthFactorEqs)
                else
                        call omegas(u, a, omega_m, omega_v)
                        growthFactor = u%growthFactorNorm * a * unnormalizedGrowthFactor(omega_m, omega_v)
                end if

        end function growthFactor

        real(D) function unnormalizedGrowthFactor(omega_m, omega_v)

                real(D), intent(in) :: omega_m, omega_v

                unnormalizedGrowthFactor = omega_m &
                    / ( omega_m ** (4.0_D / 7.0_D) - omega_v &
                    + (1.0_D + 0.5_D * omega_m) * (1.0_D + omega_v / 70.0_D) )

        end function unnormalizedGrowthFactor

        pure subroutine omegas_universe(u, a, omega_m, omega_v)

                type(tUniverse), intent(in) :: u
                real(D), intent(in) :: a
                real(D), intent(out) :: omega_m, omega_v

                omega_m = u%omega_m0 / (a * ( 1.0_D + u%omega_v0 * (a * a - 1.0_D) ) + u%omega_m0 * ( 1.0_D - a ))
                omega_v = omega_m * a ** 3 * u%omega_v0 / u%omega_m0

        end subroutine omegas_universe

        !********************************************************************
        ! This is the growing factor for delta, e.g. Peebles 1980, eq. 10.12:
        !********************************************************************
        real(D) function exactGrowthFactor(u, a)

                use mIntegration

                type(tUniverse), intent(in) :: u
                real(D), intent(in) :: a
                ! Relative delta to decide convergence of integrals:
                real(D), parameter :: EPS = 0.001_D
!                real(D), parameter :: EPS = epsilon(1.0_D)
                real(D) :: growthFactorNorm
                integer:: info
                type(tIntegral) :: growthFactor
                call new(growthFactor, 0.0_D, 1.0_D, EPS, open = .true., universe = u)
                call integrate(growthFactor, growthFactorIntegrand, growthFactorNorm)

!                if( .not. converged(growthFactor) ) then
!                        write(0, *) 'WARNING: GrowthFactorNorm integral failed'
!                        write(0, *) '         to converge.'
!                end if
                call checkConvergence(growthFactor, 'GrowthFactorNorm integral failed to converge')

                growthFactorNorm = sqrt(u%omega_m0 + u%omega_v0) & ! This is H(a = 1) / H_0
                                   * growthFactorNorm

                call new(growthFactor, 0.0_D, a, EPS, open=.TRUE., universe = u)
                call integrate(growthFactor, growthFactorIntegrand, exactGrowthFactor)

!                convergence: if( .not. converged(growthFactor) ) then
!                        write(0, *) 'WARNING: GrowthFactor integral failed to'
!                        write(0, *) '         converge.'
!                end if convergence
                call checkConvergence(growthFactor, 'GrowthFactor integral failed to converge.')

                exactGrowthFactor = sqrt(u%omega_m0 / a ** 3 + u%omega_v0) * exactGrowthFactor / growthFactorNorm

        end function exactGrowthFactor

        real(D) pure function growthFactorIntegrand(integral, a)
                use mIntegration
                type(tIntegral), intent(in) :: integral
                real(D), intent(in) :: a ! 1 / (1 + z)
                type(tUniverse) :: u

                u = getUniverse(integral)
                growthFactorIntegrand = ( a / (u%omega_m0 + u%omega_v0 * a ** 3) ) ** (1.5_D)
        end function growthFactorIntegrand

        function darkEnergyGrowthFactorDerivs(diffEq, a, yValues)
                
                use mDiffEq

                type(tDiffEq), intent(in) :: diffEq
                real(D), intent(in) :: a, yValues(:)
                real(D) :: darkEnergyGrowthFactorDerivs( size(yValues) )

                real(D) :: w, Omega_de, w0, wa, om

                call params(diffEq, om, w0, wa)
                w = w0 + (1.0_D - a) * wa

                Omega_de = om / (1.0_D - om) &
                           * exp(3.0_D * ( (wa + w0) * log(a) + wa * (1.0_D - a) ))
                Omega_de = 1.0_D / (1.0_D + Omega_de)

                darkEnergyGrowthFactorDerivs = (/ yValues(2), &
                                                       1.5_D * (w - 1.0_D) * Omega_de * yValues(1) / a ** 2 + 0.5_D * (3.0_D * w * Omega_de - 7.0_D) * yValues(2) / a /)

        end function darkEnergyGrowthFactorDerivs

        logical pure function equal(u1, u2)
                type(tUniverse), intent(in) :: u1, u2
                equal = (u1%omega_m0 == u2%omega_m0 &
                    .and. u1%omega_v0 == u2%omega_v0 &
                    .and. u1%w_0 == u2%w_0 &
                    .and. u1%w_a == u2%w_a &
                    .and. u1%omega_b0 == u2%omega_b0 &
                    .and. u1%h == u2%h &
                    .and. u1%sigma_8 == u2%sigma_8 &
                    .and. u1%b_1 == u2%b_1 &
                    .and. u1%b_2 == u2%b_2)
        end function equal

        subroutine write_universe(u)
                use mDiffEq
                type(tUniverse), intent(in) :: u
                write(6, 1) '# Omega_m0:', u%omega_m0
                write(6, 1) '# Omega_v0:', u%omega_v0
                write(6, 1) '# w_0:', u%w_0
                write(6, 1) '# w_a:', u%w_a
                write(6, 1) '# Omega_b0:', u%omega_b0
                write(6, 1) '# h:', u%h
                write(6, 1) '# sigma_8:', u%sigma_8
                write(6, 1) '# b_1:', u%b_1
                write(6, 1) '# b_2:', u%b_2
                write(6, 1) '# Shape Parameter:', getShapeParameter(u)
                write(6, 2) '# Average Density:', getAvgDensity(u), ' (M_sun/h)*(h/Mpc)**3'
1               format(a, f5.2)
2               format(a, e9.2, a)
        end subroutine write_universe

        logical pure function darkEnergy(u)
                type(tUniverse), intent(in) :: u
                darkEnergy = (u%w_0 /= -1.0_D .or. u%w_a /= 0.0_D)
        end function darkEnergy

        real(D) pure function getShapeParameter_universe(u)
                type(tUniverse), intent(in) :: u
                getShapeParameter_universe = u%omega_m0 * u%h
        end function getShapeParameter_universe

        real(D) pure function getH_universe(u)
                type(tUniverse), intent(in) :: u
                getH_universe = u%h
        end function getH_universe

        real(D) pure function getSigma_8_universe(u)
                type(tUniverse), intent(in) :: u
                getSigma_8_universe = u%sigma_8
        end function getSigma_8_universe

        real(D) pure function getOmega_m0_universe(u)
                type(tUniverse), intent(in) :: u
                getOmega_m0_universe = u%Omega_m0
        end function getOmega_m0_universe

        real(D) pure function getOmega_v0_universe(u)
                type(tUniverse), intent(in) :: u
                getOmega_v0_universe = u%Omega_v0
        end function getOmega_v0_universe

        real(D) pure function getOmega_b0_universe(u)
                type(tUniverse), intent(in) :: u
                getOmega_b0_universe = u%Omega_b0
        end function getOmega_b0_universe

        real(D) pure function getB_1_universe(u)
                type(tUniverse), intent(in) :: u
                getB_1_universe = u%b_1
        end function getB_1_universe

        real(D) pure function getB_2_universe(u)
                type(tUniverse), intent(in) :: u
                getB_2_universe = u%b_2
        end function getB_2_universe

        real(D) pure function getAvgDensity_universe(u)
                type(tUniverse), intent(in) :: u
                getAvgDensity_universe = AVG_DENSITY_OVER_H2_OMEGA_M0 * u%h * u%h * u%omega_m0
        end function getAvgDensity_universe

        real(D) function getW_0_universe(u)
                type(tUniverse), intent(in) :: u
                getW_0_universe = u%w_0
        end function getW_0_universe

        subroutine setW_0_universe(u, w_0)
                type(tUniverse), intent(inout) :: u
                real(D), intent(in) :: w_0
                u%w_0 = w_0
                call normalize(u)
        end subroutine setW_0_universe

        real(D) function getW_a_universe(u)
                type(tUniverse), intent(in) :: u
                getW_a_universe = u%w_a
        end function getW_a_universe

!        real(D) function w_0(u)
!                type(tUniverse), intent(in) :: u
!                w_0 = u%w_0
!        end function w_0

        subroutine setW_a_universe(u, w_a)
                type(tUniverse), intent(inout) :: u
                real(D), intent(in) :: w_a
                u%w_a = w_a
                call normalize(u)
        end subroutine setW_a_universe

        pure subroutine setB_1_universe(u, b_1)
                type(tUniverse), intent(inout) :: u
                real(D), intent(in) :: b_1
                u%b_1 = b_1
        end subroutine setB_1_universe

        pure subroutine setB_2_universe(u, b_2)
                type(tUniverse), intent(inout) :: u
                real(D), intent(in) :: b_2
                u%b_2 = b_2
        end subroutine setB_2_universe

        real(D) pure function getA_s_universe(u)
                type(tUniverse), intent(in) :: u
                getA_s_universe = u%A_s
        end function getA_s_universe

        pure subroutine getCMBParameters(u, omega_b, omega_d, Omega_DE)
                type(tUniverse), intent(in) :: u
                real(D), intent(out), optional :: omega_b, omega_d, Omega_DE
                if( present(omega_b) ) omega_b = u%Omega_b0 * u%h ** 2
                if( present(omega_d) ) omega_d = (u%Omega_m0 - u%Omega_b0) * u%h ** 2
                if( present(Omega_DE) ) Omega_DE = 1.0_D - u%Omega_m0
        end subroutine getCMBParameters

        subroutine setCMBParameters(u, omega_b, omega_d, Omega_DE)
                type(tUniverse), intent(inout) :: u
                real(D), intent(in) :: omega_b, omega_d, Omega_DE
                u%Omega_v0 = Omega_DE ! DEREK --- Omega_v0 is redundant, since we can't handle Omega_m0 + Omega_v0 /= 1 anyway.
                u%Omega_m0 = 1.0_D - Omega_DE
                u%Omega_b0 = omega_b / (omega_d + omega_b) * u%Omega_m0
                u%h = sqrt( (omega_d + omega_b) / u%Omega_m0 )
                call normalize(u)
        end subroutine setCMBParameters

        real(D) pure function getOmega_b_universe(u)
                type(tUniverse), intent(in) :: u
                getOmega_b_universe = u%Omega_b0 * u%h ** 2
        end function getOmega_b_universe

        subroutine setOmega_b_universe(u, omega_b)
                type(tUniverse), intent(inout) :: u
                real(D), intent(in) :: omega_b
                real(D) :: oldOmega_b, omega_d, Omega_DE
                call getCMBParameters(u, oldOmega_b, omega_d, Omega_DE)
                call setCMBParameters(u, omega_b, omega_d, Omega_DE)
        end subroutine setOmega_b_universe

        real(D) pure function getOmega_d_universe(u)
                type(tUniverse), intent(in) :: u
                getOmega_d_universe = (u%Omega_m0 - u%Omega_b0) * u%h ** 2
        end function getOmega_d_universe

        subroutine setOmega_d_universe(u, omega_d)
                type(tUniverse), intent(inout) :: u
                real(D), intent(in) :: omega_d
                real(D) :: omega_b, oldOmega_d, Omega_DE
                call getCMBParameters(u, omega_b, oldOmega_d, Omega_DE)
                call setCMBParameters(u, omega_b, omega_d, Omega_DE)
        end subroutine setOmega_d_universe

        real(D) pure function getOmega_de_universe(u)
                type(tUniverse), intent(in) :: u
                getOmega_de_universe = u%omega_v0
        end function getOmega_de_universe

        subroutine setOmega_de_universe(u, Omega_de)
                type(tUniverse), intent(inout) :: u
                real(D), intent(in) :: Omega_de
                real(D) :: omega_b, omega_d, oldOmega_de
                call getCMBParameters(u, omega_b, omega_d, oldOmega_de)
                call setCMBParameters(u, omega_b, omega_d, Omega_de)
        end subroutine setOmega_de_universe

        !******
        ! H(a)
        !******
        real(D) pure function hubble(u, a)
                type(tUniverse), intent(in) :: u
                real(D),intent(in) :: a
!                hubble = 1e5_D * u%h * sqrt(u%omega_m0 / a ** 3 + u%omega_v0)
                hubble = 1e5_D * u%h &
                        * sqrt( u%omega_m0 / a ** 3 &
                               + u%omega_v0 * exp(-3.0_D * &
                                   ( (1.0_D + u%w_0 + u%w_a) * log(a) &
                                     + u%w_a * (1.0_D - a) )) )                
        end function hubble
    
        !*********************************
        ! m as a function of Lagrangian R
        !*********************************
        real(D) pure function R_m(u, R)
                use mConstants
                type(tUniverse), intent(in) :: u
                real(D), intent(in) :: R
                R_m = 4.0_D * PI / 3.0_D * getAvgDensity(u) * R ** 3
        end function R_m

        !*********************************
        ! Lagrangian R as a function of m
        !*********************************
        real(D) pure function m_R(u, m)
                use mConstants
                type(tUniverse), intent(in) :: u
                real(D), intent(in) :: m
                m_R = ( 3.0_D * m / (4.0_D * PI * getAvgDensity(u)) ) ** ONE_THIRD
        end function m_R

        real(D) pure function z_a(z) ! Scale factor
                real(D), intent(in) :: z
                z_a = 1.0_D / (1.0_D + z)
        end function z_a

        real(D) pure function a_z(a)
                real(D), intent(in) :: a
                a_z = 1.0_D / a - 1.0_D
        end function a_z

end module mUniverse
